import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import TimeSeriesSplit
from sklearn.preprocessing import StandardScaler
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Strategy:
    def __init__(self, name):
        """
        初始化策略
        :param name: 策略名称
        """
        self.name = name
    
    def generate_signals(self, df):
        """
        生成交易信号
        :param df: 包含市场数据的DataFrame
        :return: 包含交易信号的DataFrame
        """
        raise NotImplementedError("子类必须实现此方法")

class PriceVolumeStrategy(Strategy):
    def __init__(self):
        """
        初始化量价策略
        """
        super().__init__("量价策略")
    
    def generate_signals(self, df):
        """
        基于量价关系和技术指标生成交易信号
        :param df: 包含市场数据的DataFrame
        :return: 包含交易信号的DataFrame
        """
        if df.empty:
            logger.warning("数据为空，无法生成交易信号")
            return df
        
        try:
            signals = df.copy()
            
            # 信号初始化
            signals['signal'] = 0
            
            # 买入信号：MACD金叉、KDJ金叉、RSI超卖后回升、价格突破上轨
            buy_conditions = (
                (signals['DIF'] > signals['DEA']) & 
                (signals['DIF'].shift(1) <= signals['DEA'].shift(1)) &
                (signals['K'] > signals['D']) & 
                (signals['K'].shift(1) <= signals['D'].shift(1)) &
                (signals['RSI'] < 30) & 
                (signals['RSI'] > signals['RSI'].shift(1)) &
                (signals['close'] > signals['BB_upper']) &
                (signals['volume_MA5'] > signals['volume_MA10'])
            )
            
            # 卖出信号：MACD死叉、KDJ死叉、RSI超买后回落、价格跌破下轨
            sell_conditions = (
                (signals['DIF'] < signals['DEA']) & 
                (signals['DIF'].shift(1) >= signals['DEA'].shift(1)) &
                (signals['K'] < signals['D']) & 
                (signals['K'].shift(1) >= signals['D'].shift(1)) &
                (signals['RSI'] > 70) & 
                (signals['RSI'] < signals['RSI'].shift(1)) &
                (signals['close'] < signals['BB_lower']) &
                (signals['volume_MA5'] < signals['volume_MA10'])
            )
            
            # 生成买卖信号
            signals.loc[buy_conditions, 'signal'] = 1
            signals.loc[sell_conditions, 'signal'] = -1
            
            # 生成持仓信号
            signals['position'] = signals['signal'].cumsum()
            
            logger.info(f"量价策略信号生成完成，买入信号: {len(signals[signals['signal'] == 1])}，卖出信号: {len(signals[signals['signal'] == -1])}")
            return signals
        except Exception as e:
            logger.error(f"量价策略信号生成失败: {e}")
            return df

class FinancialStrategy(Strategy):
    def __init__(self):
        """
        初始化财务指标策略
        """
        super().__init__("财务指标策略")
    
    def generate_signals(self, df):
        """
        基于财务指标生成交易信号
        :param df: 包含市场数据和财务指标的DataFrame
        :return: 包含交易信号的DataFrame
        """
        if df.empty:
            logger.warning("数据为空，无法生成交易信号")
            return df
        
        try:
            signals = df.copy()
            
            # 信号初始化
            signals['signal'] = 0
            
            # 财务指标权重
            weights = {
                'roe': 0.2,          # 净资产收益率
                'roa': 0.15,         # 总资产收益率
                'gross_margin': 0.1, # 毛利率
                'net_margin': 0.1,   # 净利率
                'revenue_grow': 0.15,# 营业收入增长率
                'profit_grow': 0.2,  # 净利润增长率
                'debt_to_assets': -0.1  # 资产负债率(负向指标)
            }
            
            # 计算综合得分
            signals['financial_score'] = 0
            for indicator, weight in weights.items():
                if indicator in signals.columns:
                    # 标准化指标值
                    signals[f'{indicator}_norm'] = (signals[indicator] - signals[indicator].mean()) / signals[indicator].std()
                    # 计算加权得分
                    signals['financial_score'] += signals[f'{indicator}_norm'] * weight
            
            # 买入信号：财务得分上升且为正，价格上涨，成交量增加
            buy_conditions = (
                (signals['financial_score'] > 0) &
                (signals['financial_score'] > signals['financial_score'].shift(1)) &
                (signals['close'] > signals['close'].shift(1)) &
                (signals['vol'] > signals['vol'].shift(1))
            )
            
            # 卖出信号：财务得分下降且为负，价格下跌，成交量减少
            sell_conditions = (
                (signals['financial_score'] < 0) &
                (signals['financial_score'] < signals['financial_score'].shift(1)) &
                (signals['close'] < signals['close'].shift(1)) &
                (signals['vol'] < signals['vol'].shift(1))
            )
            
            # 生成买卖信号
            signals.loc[buy_conditions, 'signal'] = 1
            signals.loc[sell_conditions, 'signal'] = -1
            
            # 生成持仓信号
            signals['position'] = signals['signal'].cumsum()
            
            logger.info(f"财务指标策略信号生成完成，买入信号: {len(signals[signals['signal'] == 1])}，卖出信号: {len(signals[signals['signal'] == -1])}")
            return signals
        except Exception as e:
            logger.error(f"财务指标策略信号生成失败: {e}")
            return df

class MLStrategy(Strategy):
    def __init__(self):
        """
        初始化机器学习策略
        """
        super().__init__("机器学习策略")
        self.model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.scaler = StandardScaler()
        self.feature_cols = []
    
    def train(self, df, lookback_days=30):
        """
        训练机器学习模型
        :param df: 包含市场数据和财务指标的DataFrame
        :param lookback_days: 回望天数
        """
        if df.empty:
            logger.warning("数据为空，无法训练模型")
            return
        
        try:
            # 选择特征列
            self.feature_cols = [
                'open', 'high', 'low', 'close', 'vol', 'amount', 'MA5', 'MA10', 
                'MA20', 'MA30', 'DIF', 'DEA', 'MACD', 'K', 'D', 'J', 'RSI', 
                'volume_MA5', 'volume_MA10', 'BB_MA', 'BB_std', 'BB_upper', 'BB_lower',
                'roe', 'roa', 'gross_margin', 'net_margin', 'debt_to_assets', 
                'current_ratio', 'quick_ratio', 'eps', 'bps', 'revenue_grow', 
                'profit_grow', 'assets_grow'
            ]
            
            # 确保所有特征列都存在
            self.feature_cols = [col for col in self.feature_cols if col in df.columns]
            
            if not self.feature_cols:
                logger.warning("没有可用的特征列，无法训练模型")
                return
            
            # 创建目标变量：未来5天收益率大于0标记为1，否则为0
            df['future_return'] = df['close'].pct_change(5).shift(-5)
            df['target'] = (df['future_return'] > 0).astype(int)
            
            # 移除包含NaN的行
            df = df.dropna()
            
            # 特征标准化
            X = self.scaler.fit_transform(df[self.feature_cols])
            y = df['target'].values
            
            # 使用时间序列交叉验证
            tscv = TimeSeriesSplit(n_splits=5)
            
            # 训练模型
            for train_idx, test_idx in tscv.split(X):
                X_train, X_test = X[train_idx], X[test_idx]
                y_train, y_test = y[train_idx], y[test_idx]
                
                self.model.fit(X_train, y_train)
                train_score = self.model.score(X_train, y_train)
                test_score = self.model.score(X_test, y_test)
                
                logger.info(f"模型训练完成，训练集准确率: {train_score:.4f}，测试集准确率: {test_score:.4f}")
            
            # 特征重要性
            feature_importance = pd.DataFrame({
                'feature': self.feature_cols,
                'importance': self.model.feature_importances_
            }).sort_values('importance', ascending=False)
            
            logger.info("特征重要性:\n" + str(feature_importance.head(10)))
        except Exception as e:
            logger.error(f"机器学习模型训练失败: {e}")
    
    def generate_signals(self, df):
        """
        基于机器学习模型生成交易信号
        :param df: 包含市场数据和财务指标的DataFrame
        :return: 包含交易信号的DataFrame
        """
        if df.empty or not self.feature_cols:
            logger.warning("数据为空或未训练模型，无法生成交易信号")
            return df
        
        try:
            signals = df.copy()
            
            # 信号初始化
            signals['signal'] = 0
            
            # 确保所有特征列都存在
            available_cols = [col for col in self.feature_cols if col in signals.columns]
            
            if len(available_cols) != len(self.feature_cols):
                logger.warning(f"特征列不完整，期望{len(self.feature_cols)}列，实际{len(available_cols)}列")
                missing_cols = set(self.feature_cols) - set(available_cols)
                logger.warning(f"缺失的特征列: {missing_cols}")
                return signals
            
            # 特征标准化
            X = self.scaler.transform(signals[self.feature_cols])
            
            # 预测买入概率
            signals['buy_prob'] = self.model.predict_proba(X)[:, 1]
            
            # 买入信号：预测买入概率大于阈值，且概率呈上升趋势
            buy_conditions = (
                (signals['buy_prob'] > 0.65) &
                (signals['buy_prob'] > signals['buy_prob'].shift(1)) &
                (signals['buy_prob'] > signals['buy_prob'].shift(2))
            )
            
            # 卖出信号：预测买入概率小于阈值，且概率呈下降趋势
            sell_conditions = (
                (signals['buy_prob'] < 0.45) &
                (signals['buy_prob'] < signals['buy_prob'].shift(1)) &
                (signals['buy_prob'] < signals['buy_prob'].shift(2))
            )
            
            # 生成买卖信号
            signals.loc[buy_conditions, 'signal'] = 1
            signals.loc[sell_conditions, 'signal'] = -1
            
            # 生成持仓信号
            signals['position'] = signals['signal'].cumsum()
            
            logger.info(f"机器学习策略信号生成完成，买入信号: {len(signals[signals['signal'] == 1])}，卖出信号: {len(signals[signals['signal'] == -1])}")
            return signals
        except Exception as e:
            logger.error(f"机器学习策略信号生成失败: {e}")
            return signals    